#include "Hungarian.h"
#include "lapjv.hpp"
#include <chrono>
#include <iostream>
#include <vector>

// fill matrix
double *fillMatrix(const vector<vector<double>> &costMatrix) {
  int rows = costMatrix.size();
  int cols = costMatrix[0].size();
  int dim = max(rows, cols);
  double *assign_cost = new double[dim * dim];

  for (int i = 0; i < dim; ++i) {
    for (int j = 0; j < dim; ++j) {
      if (i < rows && j < cols) {
        assign_cost[i * dim + j] = costMatrix[i][j];
      } else {
        assign_cost[i * dim + j] = 0.0;
      }
    }
  }

  return assign_cost;
}

vector<vector<double>> generateRandomMatrix(int rows, int cols) {
  vector<vector<double>> costMatrix(rows, vector<double>(cols, 0));
  for (int i = 0; i < rows; ++i) {
    for (int j = 0; j < cols; ++j) {
      costMatrix[i][j] = rand() % 50;
    }
  }
  return costMatrix;
}

int main() {
  using namespace std;
  using namespace chrono;

  vector<vector<double>> costMatrix = generateRandomMatrix(40, 30);

  // Hungarian Algorithm Test
  HungarianAlgorithm HungAlgo;
  vector<int> hungarian_assignment;
  double hungarian_total_cost = 0.0;

  // 开始计时 Hungarian
  auto hungarian_start = high_resolution_clock::now();
  for (int i = 0; i < 10000; ++i) {
    double cost = HungAlgo.Solve(costMatrix, hungarian_assignment);
    hungarian_total_cost += cost;
  }
  auto hungarian_end = high_resolution_clock::now();
  auto hungarian_duration =
      duration_cast<milliseconds>(hungarian_end - hungarian_start).count();

  // LAPJV Algorithm Test

  // 开始计时 LAPJV
  auto lapjv_start = high_resolution_clock::now();
  int dim = max(costMatrix.size(), costMatrix[0].size());
  double lapjv_total_cost = 0.0;
  vector<int> rowsol(dim), colsol(dim);
  vector<double> u(dim), v(dim);
  double *assign_cost = fillMatrix(costMatrix);

  for (int i = 0; i < 10000; ++i) {
    lapjv_total_cost += lap<false, false>(dim, assign_cost, rowsol.data(),
                                          colsol.data(), u.data(), v.data());
  }
  auto lapjv_end = high_resolution_clock::now();
  auto lapjv_duration =
      duration_cast<milliseconds>(lapjv_end - lapjv_start).count();

  // 输出结果
  cout << "Hungarian Algorithm Average Cost: " << hungarian_total_cost / 10000.0
       << endl;
  cout << "Hungarian Algorithm Total Time for 10000 runs: "
       << hungarian_duration << " ms" << endl;
  // cout hungarian_assignment
  std::cout << "Hungarian Algorithm Row assignments: ";
  for (int i = 0; i < dim; ++i) {
    std::cout << hungarian_assignment[i] << " ";
  }
  std::cout << std::endl;

  cout << "LAPJV Algorithm Average Cost: " << lapjv_total_cost / 10000.0
       << endl;
  cout << "LAPJV Algorithm Total Time for 10000 runs: " << lapjv_duration
       << " ms" << endl;
  std::cout << "Row assignments: ";
  for (int i = 0; i < dim; ++i) {
    std::cout << rowsol[i] << " ";
  }
  std::cout << std::endl;

  std::cout << "Column assignments: ";
  for (int j = 0; j < dim; ++j) {
    std::cout << colsol[j] << " ";
  }
  std::cout << std::endl;

  delete[] assign_cost;

  return 0;
}
